# This code provides the values used for Figure 5
# RStan is required

rm(list=ls())
library(rstan)
library(ggplot2)
library(lattice)
setwd("")
data <- read.dta("data.dta")

## Function for Stan BIQQ model (see background material to Humphreys and Jacobs 2015)
biqq.stan.code <- "

# User supplies the following data:
data {
  # Counts of the different XY and XYK outcomes observed:
  int<lower=0> XYK[8]; # summary of N of XYK outcomes
  int<lower=0> XY[4]; # summary of N of XY outcomes
  # And the hyperparameters for the prior:
  vector[4] alpha_prior; # Dirichlet shape parameters for share of types (abcd)
  vector[4] pi_alpha; # Beta shape parameters for assignment probabilities
  vector[4] pi_beta ;
  vector[4] phi0_alpha ; # Beta shape parameters for probative value of clues
  vector[4] phi0_beta ;
  vector[4] phi1_alpha ;
  vector[4] phi1_beta ;
}
# The parameters to be modelled are defined:
parameters {
  simplex[4] abcd; # The share of abcd types in the population
  # (constrained to sum to 1)
  real<lower=0,upper=1> pi_a; # The probabilities of assignment
  real<lower=0,upper=1> pi_b; # (constrained between 0 and 1 inclusive)
  real<lower=0,upper=1> pi_c;
  real<lower=0,upper=1> pi_d;
  real<lower=0,upper=1> phi_a0; # The probative values of clues for each type
  real<lower=0,upper=1> phi_b0; # (constrained between 0 and 1 inclusive)
  real<lower=0,upper=1> phi_c0;
  real<lower=0,upper=1> phi_d0;
  real<lower=0,upper=1> phi_a1;
  real<lower=0,upper=1> phi_b1;
  real<lower=0,upper=1> phi_c1;
  real<lower=0,upper=1> phi_d1;
}
# The multinomial event probabilities are calculated:
transformed parameters {
  simplex[4] w_XY;
  simplex[8] w_XYK;
  w_XY[1] <- (1-pi_b)*abcd[2] + (1-pi_c)*abcd[3] ; # Pr(00*)
  w_XY[2] <- (1-pi_a)*abcd[1] + (1-pi_d)*abcd[4] ; # Pr(01*)
  w_XY[3] <- pi_a*abcd[1] + pi_c*abcd[3] ; # Pr(10*)
  w_XY[4] <- pi_b*abcd[2] + pi_d*abcd[4] ; # Pr(11*)
  w_XYK[1]<- (1-pi_b)*(1-phi_b0)*abcd[2] + (1-pi_c)*(1-phi_c0)*abcd[3]; # Pr(000)
  w_XYK[2]<- (1-pi_b)*phi_b0*abcd[2] + (1-pi_c)*phi_c0*abcd[3] ; # Pr(001)
  w_XYK[3]<- (1-pi_a)*(1-phi_a0)*abcd[1] + (1-pi_d)*(1-phi_d0)*abcd[4]; # Pr(010)
  w_XYK[4]<- (1-pi_a)*phi_a0*abcd[1] + (1-pi_d)*phi_d0*abcd[4] ; # Pr(011)
  w_XYK[5]<- pi_a*(1-phi_a1)*abcd[1] + pi_c*(1-phi_c1)*abcd[3] ; # Pr(100)
  w_XYK[6]<- pi_a*phi_a1*abcd[1] + pi_c*phi_c1*abcd[3] ; # Pr(101)
  w_XYK[7]<- pi_b*(1-phi_b1)*abcd[2] + pi_d*(1-phi_d1)*abcd[4] ; # Pr(110)
  w_XYK[8]<- pi_b*phi_b1*abcd[2] + pi_d*phi_d1*abcd[4] ; # Pr(111)
}
# The parameters are then modeled as follows:
model {
  abcd ~ dirichlet(alpha_prior); # Priors for abcd types
  pi_a ~ beta(pi_alpha[1], pi_beta[1]); # Priors for assignment probs
  pi_b ~ beta(pi_alpha[2], pi_beta[2]);
  pi_c ~ beta(pi_alpha[3], pi_beta[3]);
  pi_d ~ beta(pi_alpha[4], pi_beta[4]);
  phi_a0 ~ beta(phi0_alpha[1], phi0_beta[1]); # Priors for clues
  phi_a1 ~ beta(phi1_alpha[1], phi1_beta[1]);
  phi_b0 ~ beta(phi0_alpha[2], phi0_beta[2]);
  phi_b1 ~ beta(phi1_alpha[2], phi1_beta[2]);
  phi_c0 ~ beta(phi0_alpha[3], phi0_beta[3]);
  phi_c1 ~ beta(phi1_alpha[3], phi1_beta[3]);
  phi_d0 ~ beta(phi0_alpha[4], phi0_beta[4]);
  phi_d1 ~ beta(phi1_alpha[4], phi1_beta[4]);
  XY ~ multinomial(w_XY); # Likelihood Part 1
  XYK ~ multinomial(w_XYK); # Likelihood Part 2
}
"


## Input hypothetical cross-case data to start model (again, code provided in SI to Humphreys and Jacobs (2015))

# Compile the model first, giving it simple starting values
ones <- rep(1,4)
biqq.stan.model <- stan(
  # The code for the model
  model_code = biqq.stan.code,
  # The initial data
  data = list(XY = c( 1, # Number of 00*
                      1, # Number of 01*
                      1, # Number of 10*
                      1 # Number of 11*
  ),
  XYK = c( 1, 1, # Number of 000 & 001
           1, 1, # Number of 010 & 011
           1, 1, # Number of 100 & 101
           1, 1 # Number of 110 & 111
  ),
  # Shape parameters for flat priors:
  alpha_prior = ones,
  pi_alpha = ones, pi_beta = ones,
  phi0_alpha = ones, phi0_beta = ones,
  phi1_alpha = ones, phi1_beta = ones),
  # Run a low number of iterations to warm the model up:
  iter = 100,
  chains = 4
)


# BIQQ function for drawing from posterior distribution using STAN:
biqq <- function(fit = biqq.stan.model,
                 # Define the default arguments of the function:
                 XY = c( 1, 1, 1, 1),
                 XYK = c( 1, 1,
                          1, 1,
                          1, 1,
                          1, 1),
                 alpha_prior = ones,
                 pi_alpha = ones, pi_beta = ones,
                 phi0_alpha = ones, phi0_beta = ones,
                 phi1_alpha = ones, phi1_beta = ones,
                 iter = 1000, chains = 4,
                 warmup = 100, seed = 100
) {
  # Function takes inputs on data and priors and implements the baseline BIQQ model using rstan
  # Define the model inputs, based on user-supplied or default arguments
  data <- list(
    XY = XY,
    XYK = XYK,
    alpha_prior = alpha_prior,
    pi_alpha = pi_alpha, pi_beta = pi_beta,
    phi0_alpha = phi0_alpha, phi0_beta = phi0_beta,
    phi1_alpha = phi1_alpha, phi1_beta = phi1_beta
  )
  # Use Stan to sample from the posterior distribution:
  posterior <-
    stan(fit = fit,
         data = data,
         iter = iter,
         chains = chains,
         warmup = warmup,
         seed = seed
    )
  # Display the results
  return(posterior)
}
# Demonstration 1:
# Simple demonstration using default values
biqq()


## Put in priors from experts

# Here, we paste the priors from the experts.
# They indicate whether for a given type (X=Assignment, Y=Outcome),
# a clue (K=Clue) yielded positive causal evidence (1) or not (0)

XYK_obedience =  c(2 ,   1,   0,   0,   1,   4,   0,   8)
XYK_persuasion = c(0 ,   0,   0,   0,   0,   6,   0,   7)

XY_both_empty <- rep(0, 4)


## Assume flat priors on alpha, phi and pi

ap <-  c(1, 1, 1, 1) # prior for alpha -- 1 is flat (0 not allowed)
pip <- c(1, 1, 1, 1) # prior for pi -- 1 is flat (0 not allowed)

b1 <- biqq(XYK = XYK_obedience, XY = XY_both_empty, alpha_prior = ap, pi_alpha = pip) # obedience
b2 <- biqq(XYK = XYK_persuasion, XY = XY_both_empty, alpha_prior = ap, pi_alpha = pip) # persuasion

## function to get stats, make table, print

stats = function(x1,x2){c(
  mean.ab  = mean(x1),
  sd.ab    = sd(x1),
  lower.ab = quantile(x1,probs = .025),
  upper.ab = quantile(x1,probs = 1-.025),
  mean.c   = mean(x2),
  sd.c     = sd(x2))
}

## get table row function

row.makr <- function(out){
  d <- extract(out)  
  stats(d$abcd[,2]-d$abcd[,1],d$pi_c)
}

## get exemplary table rows

row.1 <- row.makr(b1)  
row.2 <- row.makr(b2) 

## join rows, name rows, make table

tab <- round(rbind(row.1,row.2),3)
row.names(tab) <- c("obedience", "persuasion")
tab <- cbind(tab, 20)
tab <- cbind(tab, 5)
colnames(tab)[7] <- "lambda_b"
colnames(tab)[8] <- "pi_b"

tab

## loop through ratio values of lambda_b

# get values, make plot

values_a = 1:20 # sensitivity for alpha
values_p = 1:20 # sensitivity for pi 

for (i in values_a) {
  for (j in values_p) {
  
    ## get alpha prior vector
    
    #lambda_b <- (3*i)/(1-i)
    lambda_b = i
    ap <- c(1, lambda_b, 1, 1) # alpha prior vector
    
    ## get pi prior vector
    
    pi_b <- j
    pip <- c(1, pi_b, 1, 1) # prior for pi, zero not allowed, arbitrary
    
    ## get BIQQ estimates
    
    b1 <- biqq(XYK = XYK_obedience, XY = XY_both_empty, 
               alpha_prior = ap, pi_alpha = pip) 
    b2 <- biqq(XYK = XYK_persuasion, XY = XY_both_empty, 
               alpha_prior = ap, pi_alpha = pip) 
    
    # ratio to table row, to df, bind to table
    
    lambda_b 
    row.1 <- c(row.makr(b1), lambda_b, pi_b) 
    row.2 <- c(row.makr(b2), lambda_b, pi_b) 
    tab <- rbind(tab, row.1, row.2)
  }
}

# Rearrange results
tab <- tab[-(1:2), ] ## remove example ratios
tab <- data.frame(tab) ## to df, ignore warnings
tab$var <- rep(c("Obedience", "Persuasion"), max(values_a) * max(values_p)) ## add outcome names

# Save results 
write.csv(tab, "bayes_results.csv", row.names = F)

## Plot results

## Insheet BIQQ results
tab <- read.csv("bayes_results.csv")
tab_ob <- tab[tab$var == "Obedience", ]
tab_pe <- tab[tab$var == "Persuasion", ]


## Coefplot

tab_plot <- tab[tab$pi_b == 1, ]

pd <- position_dodge(0.2)

p1 <- ggplot(tab_plot, aes(x = lambda_b, y = mean.ab)) + geom_point(size = 2)
p1 <- p1 + geom_linerange(aes(y = mean.ab, ymin = lower.ab.2.5., ymax = upper.ab.97.5.))
p1 <- p1 + facet_wrap(~var) + theme_bw() + xlab(expression(paste("Prior for ", lambda[b])))
p1 <- p1 + ylab("Treatment Effect") + geom_hline(yintercept = 0, linetype = "dashed", size = 0.7)
p1 <- p1 + theme(axis.line = element_line(colour = "black"),
                 panel.grid.major = element_blank(),
                 panel.grid.minor = element_blank(),
                 panel.background = element_blank()) 
p1

ggsave("BIQQ1.pdf", width = 7, height = 4)


